use alloc::vec::Vec;
use alloc::vec;
use bitflags::*;

use super::address::*;
use super::frame_allocator::*;

bitflags! {
    pub struct PTEFlags: u8 {
        const V = 1 << 0;
        const R = 1 << 1;
        const W = 1 << 2;
        const X = 1 << 3;
        const U = 1 << 4;
        const G = 1 << 5;
        const A = 1 << 6;
        const D = 1 << 7;
    }
}

#[derive(Copy, Clone)]
#[repr(C)]
pub struct PageTableEntry {
    pub bits: usize,
}

// 页表项: [63:54]: 10 个保留位, [53:10] 44 位长的 ppn, [9:8] RSW [7:0] 8 个标志位
impl PageTableEntry { // 页表项 PTE
    pub fn new(ppn: PhysPageNum, flags: PTEFlags) -> Self {
        PageTableEntry {
            bits: ppn.0 << 10 | flags.bits as usize, // 详见: file:///./sv39-pte.png
        }
    }
    pub fn empty() -> Self {
        PageTableEntry {
            bits: 0,
        }
    }
    pub fn ppn(&self) -> PhysPageNum {
        (self.bits >> 10 & ((1usize << 44) - 1)).into()
    }
    pub fn flags(&self) -> PTEFlags {
        PTEFlags::from_bits(self.bits as u8).unwrap()
    }
    pub fn is_valid(&self) -> bool {
        (self.flags() & PTEFlags::V) != PTEFlags::empty() // 当 V 为 1 时, 页表项才合法
    }
    pub fn readable(&self) -> bool {
        (self.flags() & PTEFlags::R) != PTEFlags::empty()
    }
    pub fn writable(&self) -> bool {
        (self.flags() & PTEFlags::W) != PTEFlags::empty()
    }
    pub fn executable(&self) -> bool {
        (self.flags() & PTEFlags::X) != PTEFlags::empty()
    }
}

// SV39 多级页表以节点为单位进行管理的
// 每个节点恰好存储在一个物理页帧中
// 每个非叶节点中都有 512 个页表项
// 它的位置可以用一个物理页号来表示
// 即: ppn <--> 物理页帧 <--> 多级页表节点

// 每个应用的地址空间都对应一个不同的多级页表
// 不同页表的起始地址（即页表根节点的地址）是不一样的
// 通过保存根节点的物理页号 root_ppn 作为页表唯一的区分标志
pub struct PageTable {
    // 根节点
    root_ppn: PhysPageNum,
    // 以 FrameTracker 的形式保存了页表所有的节点所在的物理页帧。
    // FrameTracker 实现了 RAII, 所以当 PageTable 生命周期结束的时候, 内存也会被自动回收 
    frames: Vec<FrameTracker>, 
}

impl PageTable {
    pub fn new() -> Self {
        let frame = frame_alloc().unwrap();
        PageTable {
            root_ppn: frame.ppn,
            frames: vec![frame],
        }
    }
}

// 多级页表并不是被创建出来之后就不再变化的
// 为了 MMU 能够通过地址转换正确找到应用地址空间中的数据实际被内核放在内存中位置
// 操作系统需要动态维护一个虚拟页号到页表项的映射，支持插入/删除键值对
// 具体来说: 需要提前扩充多级页表维护的映射
// 使得对于每一个对应于某一特定物理页帧的物理页号 ppn ，均存在一个虚拟页号 vpn 能够映射到它
// 这个映射关系, 暂时定为最简单的恒等映射: 对于物理内存上的每个物理页帧，都在多级页表中用一个与其物理页号相等的虚拟页号映射到它
impl PageTable {
    fn find_pte_create(&mut self, vpn: VirtPageNum) -> Option<&mut PageTableEntry> {
        let idxs = vpn.indexes(); // 三级索引
        let mut ppn = self.root_ppn;
        let mut result: Option<&mut PageTableEntry> = None;
        for i in 0..3 { // 三级索引
            let pte = &mut ppn.get_pte_array()[idxs[i]];
            if i == 2 { // 已经到达叶节点
                result = Some(pte); // 将结果返回
                break;
            }
            // 如果 PTE 不合法, 就说明这个 PTE 对应的节点还没有被创建
            // 就要创建新节点
            if !pte.is_valid() {
                // 申请一个 物理帧 <--> 一个多级页表节点
                let frame = frame_alloc().unwrap(); // TODO 增加对物理帧耗尽的处理
                // 创建新的页表项, 并将 V 位置为 1 
                *pte =  PageTableEntry::new(frame.ppn, PTEFlags::V);
                // 将其加入当前节点的 frames 中 方便管理物理帧的生命周期
                self.frames.push(frame);
            }
            ppn = pte.ppn();
        }
        result
    }

    // 用于 新增/修改 虚拟地址 -> 物理地址 的映射
    pub fn map(&mut self, vpn: VirtPageNum, ppn: PhysPageNum, flags: PTEFlags) {
        let pte = self.find_pte_create(vpn).unwrap(); // TODO 增加对物理帧耗尽的处理
        assert!(!pte.is_valid(), "vpn {:?} is mapped before mapping", vpn);
        *pte = PageTableEntry::new(ppn, flags | PTEFlags::V);
    }

    // 用于 删除 虚拟地址 -> 物理地址 的映射
    pub fn unmap(&mut self, vpn: VirtPageNum) {
        let pte = self.find_pte_create(vpn).unwrap(); // TODO 增加对物理帧耗尽的处理
        assert!(pte.is_valid(), "vpn {:?} is invalid before unmapping", vpn);
        *pte = PageTableEntry::empty();
    }
}


// 为 PageTable 增加不经过 MMU 而是 手动查表 的方式
impl PageTable {
    // 创建一个专用来手动查页表的 临时 PageTable
    // 它仅有一个从传入的 satp token 中得到的多级页表根节点的物理页号
    // 它的 frames 字段为空，也即不实际控制任何资源
    pub fn from_token(satp: usize) -> Self {
        Self {
            root_ppn: PhysPageNum::from(satp & ((1usize) << 44) - 1),
            frames: Vec::new(),
        }
    }

    pub fn find_vpn(&self, vpn: VirtPageNum) -> bool {
        match self.find_pte(vpn) {
            None => false,
            Some(x) => x.is_valid(),
        }
    }

    // 只是查表, 不会试图分配物理页帧
    fn find_pte(&self, vpn: VirtPageNum) -> Option<&PageTableEntry> {
        let idxs = vpn.indexes();
        let mut ppn = self.root_ppn;
        let mut result: Option<&PageTableEntry> = None;
        for i in 0..3 {
            let pte = &ppn.get_pte_array()[idxs[i]];
            if i == 2 {
                result = Some(pte);
                break;
            }
            if !pte.is_valid() { // PTE 非法, 直接返回 None
                return None;
            }
            ppn = pte.ppn();
        }
        result
    }

    // 若能够找到页表项，则将页表项拷贝一份并返回，否则就 返回一个 None
    // 用于将 vpn 转换为 一个 PTE
    pub fn translate(&self, vpn: VirtPageNum) -> Option<PageTableEntry> {
        self.find_pte(vpn)
            .map(|pte| { pte.clone() })
    }

    // 按照 satp CSR 格式要求 构造一个无符号 64 位无符号整数
    // 使得其 分页模式为 SV39 ，且将当前多级页表的根节点所在的物理页号填充进去。
    pub fn token(&self) -> usize {
        8usize << 60 | self.root_ppn.0
    }
}

pub fn translated_byte_buffer(token: usize, ptr: *const u8, len: usize) -> Vec<&'static mut [u8]> {
    let page_table = PageTable::from_token(token);
    let mut start = ptr as usize;
    let end = start + len;
    let mut v = Vec::new();
    while start < end {
        let start_va = VirtAddr::from(start);
        let mut vpn = start_va.floor();
        let ppn = page_table
            .translate(vpn)
            .unwrap()
            .ppn();
        vpn.step();
        let mut end_va: VirtAddr = vpn.into();
        end_va = end_va.min(VirtAddr::from(end));
        if end_va.page_offset() == 0 {
            v.push(&mut ppn.get_bytes_array()[start_va.page_offset()..]);
        } else {
            v.push(&mut ppn.get_bytes_array()[start_va.page_offset()..end_va.page_offset()]);
        }
        start = end_va.into();
    }
    v
}